{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 26,
   "source": [
    "import pandas as pd\r\n",
    "import matplotlib.pyplot as plt\r\n",
    "import numpy as np\r\n",
    "from osgeo import gdal\r\n",
    "import geopandas as gpd\r\n",
    "gdal.AllRegister()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "source": [
    "gdf = gpd.read_file(\"../shp/river_pts.shp\")\r\n",
    "gdf.head()"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "   ARCID  GRID_CODE  FROM_NODE  TO_NODE     LENGTH  BASINID  TYPE  ORIG_FID  \\\n",
       "0      1          4          6        7   5.000000        4     1         0   \n",
       "1      1          4          6        7   5.000000        4     1         0   \n",
       "2      2          3          4        8  55.784271        3     1         1   \n",
       "3      2          3          4        8  55.784271        3     1         1   \n",
       "4      2          3          4        8  55.784271        3     1         1   \n",
       "\n",
       "                    geometry  \n",
       "0  POINT (6107.818 6238.657)  \n",
       "1  POINT (6112.818 6238.657)  \n",
       "2  POINT (4327.818 6253.657)  \n",
       "3  POINT (4320.318 6253.657)  \n",
       "4  POINT (4315.318 6248.657)  "
      ],
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>ARCID</th>\n",
       "      <th>GRID_CODE</th>\n",
       "      <th>FROM_NODE</th>\n",
       "      <th>TO_NODE</th>\n",
       "      <th>LENGTH</th>\n",
       "      <th>BASINID</th>\n",
       "      <th>TYPE</th>\n",
       "      <th>ORIG_FID</th>\n",
       "      <th>geometry</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>4</td>\n",
       "      <td>6</td>\n",
       "      <td>7</td>\n",
       "      <td>5.000000</td>\n",
       "      <td>4</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>POINT (6107.818 6238.657)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>4</td>\n",
       "      <td>6</td>\n",
       "      <td>7</td>\n",
       "      <td>5.000000</td>\n",
       "      <td>4</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>POINT (6112.818 6238.657)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>3</td>\n",
       "      <td>4</td>\n",
       "      <td>8</td>\n",
       "      <td>55.784271</td>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>POINT (4327.818 6253.657)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2</td>\n",
       "      <td>3</td>\n",
       "      <td>4</td>\n",
       "      <td>8</td>\n",
       "      <td>55.784271</td>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>POINT (4320.318 6253.657)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2</td>\n",
       "      <td>3</td>\n",
       "      <td>4</td>\n",
       "      <td>8</td>\n",
       "      <td>55.784271</td>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>POINT (4315.318 6248.657)</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ]
     },
     "metadata": {},
     "execution_count": 27
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "假设污染从5543.66,4810.32开始"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "source": [
    "# 读取TIFF遥感影像\r\n",
    "def read_img(filename):\r\n",
    "    dataset = gdal.Open(filename)\r\n",
    "    ncols = dataset.RasterXSize  # 栅格矩阵的列数\r\n",
    "    nrows = dataset.RasterYSize  # 栅格矩阵的行数\r\n",
    "    proj = dataset.GetProjection()  # 地图投影信息，字符串表示\r\n",
    "    data = dataset.ReadAsArray(0, 0)\r\n",
    "    datatype = data.dtype\r\n",
    "    adfGeoTransform = dataset.GetGeoTransform()\r\n",
    "    #DEM的平面四至\r\n",
    "    Xmin = adfGeoTransform[0]\r\n",
    "    Ymin = adfGeoTransform[3]\r\n",
    "    Xmax = adfGeoTransform[0] + nrows * adfGeoTransform[1] \\\r\n",
    "        + ncols * adfGeoTransform[2]\r\n",
    "    Ymax = adfGeoTransform[3] + nrows * adfGeoTransform[4] \\\r\n",
    "        + ncols * adfGeoTransform[5]\r\n",
    "    xcellwidth = abs(adfGeoTransform[1])\r\n",
    "    ycellwidth = abs(adfGeoTransform[5])\r\n",
    "    del dataset  # 关闭对象dataset，释放内存\r\n",
    "    return {\"nrows\": nrows, \"ncols\": ncols, \"data\": data,\r\n",
    "            \"Xmin\": Xmin, \"Ymin\": Ymin, \"Xmax\": Xmax, \"Ymax\": Ymax,\r\n",
    "            \"xcellwidth\": xcellwidth, \"ycellwidth\": ycellwidth,\r\n",
    "            \"proj\": proj, \"datatype\": datatype}\r\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "source": [
    "img = read_img('../raster/river.tif')\r\n",
    "img"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "{'nrows': 251,\n",
       " 'ncols': 291,\n",
       " 'data': array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "        [0., 0., 0., ..., 0., 0., 0.],\n",
       "        [0., 0., 0., ..., 0., 0., 0.],\n",
       "        ...,\n",
       "        [0., 0., 0., ..., 0., 0., 0.],\n",
       "        [0., 0., 0., ..., 0., 0., 0.],\n",
       "        [0., 0., 0., ..., 0., 0., 0.]], dtype=float32),\n",
       " 'Xmin': 4047.8177688837,\n",
       " 'Ymin': 6378.257143892,\n",
       " 'Xmax': 5829.917768883701,\n",
       " 'Ymax': 4312.157143892,\n",
       " 'xcellwidth': 7.100000000000003,\n",
       " 'ycellwidth': 7.100000000000001,\n",
       " 'proj': '',\n",
       " 'datatype': dtype('float32')}"
      ]
     },
     "metadata": {},
     "execution_count": 29
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "source": [
    "img_data = img[\"data\"]\r\n",
    "startRow = 210\r\n",
    "startCol = 110"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "source": [
    "def findLowDirection(rawdata,x,y):\r\n",
    "    \"\"\"找到某一单元周围八方向的最小值的方向\r\n",
    "\r\n",
    "    Args:\r\n",
    "        rawdata ([type]): 原始栅格数据\r\n",
    "        x ([type]): [description]\r\n",
    "        y ([type]): [description]\r\n",
    "\r\n",
    "    Returns:\r\n",
    "        [type]: [description]\r\n",
    "    \"\"\"\r\n",
    "    window = rawdata[x-1:x+2, y-1:y+2]\r\n",
    "    minValue = window[1,1]\r\n",
    "    minX = 1\r\n",
    "    minY = 1\r\n",
    "    for i in range(0,3):#row\r\n",
    "        for j in range(0,3):#col\r\n",
    "           if window[i,j] <= minValue and window[i,j]>0:\r\n",
    "             minValue = window[i,j]\r\n",
    "             minX = j\r\n",
    "             minY = i\r\n",
    "    dirc = (minX-1,minY-1) \r\n",
    "    print('dirc',dirc)\r\n",
    "    return  dirc     \r\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "source": [
    "import math\r\n",
    "def calDis(xcellwidth,ycellwidth,s_X, s_Y,e_X, e_Y):\r\n",
    "    return math.sqrt(math.pow((s_X-e_X)*xcellwidth,2)+math.pow((s_Y-e_Y)*ycellwidth,2))"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "source": [
    "def walk(resData,x, y,direction):\r\n",
    "  if direction[0] == 0 and direction[1] ==0:\r\n",
    "    return 0\r\n",
    "  else:\r\n",
    "    distance = calDis(\r\n",
    "        img['xcellwidth'], img['ycellwidth'], startRow, startCol, x+direction[0], y+direction[1])\r\n",
    "    print ('dis',distance)\r\n",
    "    newX = x+direction[1]\r\n",
    "    newY = y+direction[0]\r\n",
    "    print ('new',newX,newY)\r\n",
    "    resData[newX][newY] = distance\r\n",
    "    return 1\r\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "source": [
    "resultData = np.zeros([img[\"nrows\"], img[\"ncols\"]], dtype=float)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "X = startRow\r\n",
    "Y = startCol\r\n",
    "flag = 1\r\n",
    "while flag!= 0:\r\n",
    "    Direction = findLowDirection(img[\"data\"], X, Y)\r\n",
    "    flag=walk(resultData,X,Y,Direction)\r\n",
    "    X = X+Direction[1]\r\n",
    "    Y = Y+Direction[0]\r\n",
    "    print('loc',X,Y)\r\n",
    "    print(flag)\r\n"
   ],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "orig_nbformat": 4,
  "language_info": {
   "name": "python",
   "version": "3.8.10",
   "mimetype": "text/x-python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python",
   "file_extension": ".py"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.10 64-bit ('geo-env': conda)"
  },
  "interpreter": {
   "hash": "7b0d6b4aa929793d88e84d754c187056723d52cecd68845c7dca810866802715"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}